import torch
import os
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
from nets.myModel import models
from MyDataloader.myConfig import Config
from nets.model_layers import PriorBox

class Models(nn.Module):
    def __init__(self, phase, baseModel, num_classes, loc_layers, conf_layers):
        super(Models, self).__init__()
        self.phase = phase
        self.num_classes = num_classes
        self.model = nn.ModuleList(baseModel)
        self.loc = nn.ModuleList(loc_layers)
        self.conf = nn.ModuleList(conf_layers)
        self.priorbox = PriorBox(Config)
        with torch.no_grad():
            self.priors = Variable(self.priorbox.forward())
    def forward(self, x):
        sources = list()
        loc = list()
        conf = list()
        # 取出5x5的网格放入列表
        for k in range(44):
            x = self.model[k](x)
        sources.append(x)
        # 取出3x3的网格放入列表
        for k in range(44, len(self.model)):
            x = self.model[k](x)
        sources.append(x)
        # 利用卷积代替全连接计算先验框的回归和分类回归
        for (x, l, c) in zip(sources, self.loc, self.conf):
            loc.append(l(x).permute(0, 2, 3, 1).contiguous())
            conf.append(c(x).permute(0, 2, 3, 1).contiguous())
        # 进行resize
        loc = torch.cat([o.view(o.size(0), -1) for o in loc], 1)
        conf = torch.cat([o.view(o.size(0), -1) for o in conf], 1)
        # 返回输出，包括先验框的回归，分类回归，实际划分的先验框坐标
        output = (
            loc.view(loc.size(0), -1, 4),
            conf.view(conf.size(0), -1, self.num_classes),
            self.priors
        )

        return output


def get_model(phase, num_classes):
    loc_layers = []
    conf_layers = []
    for i in range(2):
        loc_layers += [nn.Conv2d(512, 16, kernel_size = 3, padding=1)]
        conf_layers += [nn.Conv2d(512, 8, kernel_size = 3, padding=1)]
    the_model = Models(phase, models(3), num_classes, loc_layers, conf_layers)

    return the_model